{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 测试 Pass\n",
    "\n",
    "## split args"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tvm.ir.transform import Pass\n",
    "from tvm.ir import IRModule\n",
    "from tvm.relay import transform\n",
    "from tvm import relay\n",
    "\n",
    "def run_opt_pass(expr, opt_pass):\n",
    "    assert isinstance(opt_pass, Pass)\n",
    "    mod = IRModule.from_expr(expr)\n",
    "    mod = relay.transform.InferType()(mod)\n",
    "    mod = opt_pass(mod)\n",
    "    entry = mod[\"main\"]\n",
    "    return entry if isinstance(expr, relay.Function) else entry.body"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "target = tvm.target.Target(\"metal\")\n",
    "\n",
    "shape = (1, 1, 1, 3)\n",
    "dtype = \"float32\"\n",
    "axis = 1\n",
    "inputs = []\n",
    "for i in range(100):\n",
    "    inputs.append(relay.var(f\"p{i}\", shape=shape, dtype=dtype))\n",
    "\n",
    "def before():\n",
    "    inp = relay.Tuple(inputs)\n",
    "    return relay.op.concatenate(inp, axis)\n",
    "\n",
    "res = run_opt_pass(before(), transform.SplitArgs(target.max_function_args))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "limit = target.max_function_args - 1  # one buffer with output\n",
    "splitNum = int(len(inputs) / limit)\n",
    "if len(inputs) % limit > 0:\n",
    "    splitNum += 1\n",
    "\n",
    "splitted = []\n",
    "for i in range(splitNum):\n",
    "    startIdx = i * limit\n",
    "    argsCount = min(limit, len(inputs) - startIdx)\n",
    "    args = []\n",
    "    for j in range(argsCount):\n",
    "        args.append(inputs[j + startIdx])\n",
    "    t = relay.Tuple(args)\n",
    "    concat = relay.op.concatenate(t, axis)\n",
    "    splitted.append(relay.annotation.stop_fusion(concat))\n",
    "inp = relay.Tuple(splitted)\n",
    "expr = relay.op.concatenate(inp, axis)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tvm.ir.structural_equal(res, expr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pass Instrument"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tvm\n",
    "from tvm.ir import IRModule\n",
    "from tvm import relay\n",
    "from tvm.relay import op\n",
    "from tvm import transform\n",
    "from tvm.ir.instrument import PassTimingInstrument, pass_instrument\n",
    "from tvm.ir.transform import PassContext\n",
    "\n",
    "def get_test_model():\n",
    "    x, y, z = [relay.var(c, shape=(3, 4), dtype=\"float32\") for c in \"xyz\"]\n",
    "    e1 = op.add(x, y)\n",
    "    e2 = op.subtract(x, z)\n",
    "    e3 = op.multiply(e1, e1 / e2)\n",
    "    return IRModule.from_expr(e3 + e2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pass_timing = PassTimingInstrument()\n",
    "\n",
    "seq = transform.Sequential([relay.transform.AnnotateSpans(),\n",
    "                            relay.transform.ToANormalForm(),\n",
    "                            relay.transform.InferType()])\n",
    "\n",
    "# 覆盖当前 PassContext 的 instruments\n",
    "PassContext.current().override_instruments([pass_timing])\n",
    "\n",
    "mod = get_test_model()\n",
    "mod = seq(mod)\n",
    "profiles = pass_timing.render()\n",
    "assert \"AnnotateSpans\" in profiles\n",
    "assert \"ToANormalForm\" in profiles\n",
    "assert \"InferType\" in profiles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 重置当前 PassContext 的 instruments 为 None\n",
    "PassContext.current().override_instruments(None)\n",
    "mod = get_test_model()\n",
    "mod = seq(mod)\n",
    "profiles = pass_timing.render()\n",
    "assert profiles == \"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "instrument_definition_type = tvm.testing.parameter(\"decorator\", \"subclass\")\n",
    "\n",
    "def test_custom_instrument(instrument_definition_type):\n",
    "    class BaseTest:\n",
    "        def __init__(self):\n",
    "            self.events = []\n",
    "\n",
    "        def enter_pass_ctx(self):\n",
    "            self.events.append(\"enter ctx\")\n",
    "\n",
    "        def exit_pass_ctx(self):\n",
    "            self.events.append(\"exit ctx\")\n",
    "\n",
    "        def run_before_pass(self, mod, info):\n",
    "            self.events.append(\"run before \" + info.name)\n",
    "\n",
    "        def run_after_pass(self, mod, info):\n",
    "            self.events.append(\"run after \" + info.name)\n",
    "\n",
    "    if instrument_definition_type == \"decorator\":\n",
    "        MyTest = pass_instrument(BaseTest)\n",
    "\n",
    "    elif instrument_definition_type == \"subclass\":\n",
    "        class MyTest(BaseTest, tvm.ir.instrument.PassInstrument):\n",
    "            def __init__(self):\n",
    "                BaseTest.__init__(self)\n",
    "                tvm.ir.instrument.PassInstrument.__init__(self)\n",
    "\n",
    "    mod = get_test_model()\n",
    "    my_test = MyTest()\n",
    "    with tvm.transform.PassContext(instruments=[my_test]):\n",
    "        mod = tvm.relay.transform.InferType()(mod)\n",
    "\n",
    "    assert (\n",
    "        \"enter ctx\"\n",
    "        \"run before InferType\"\n",
    "        \"run after InferType\"\n",
    "        \"exit ctx\" == \"\".join(my_test.events)\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 禁用 pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@pass_instrument\n",
    "class CustomPI:\n",
    "    def __init__(self):\n",
    "        self.events = []\n",
    "\n",
    "    def should_run(self, mod, info):\n",
    "        # Only run pass name contains \"InferType\"\n",
    "        if \"InferType\" not in info.name:\n",
    "            return False\n",
    "        return True\n",
    "\n",
    "    def run_before_pass(self, mod, info):\n",
    "        self.events.append(info.name)\n",
    "\n",
    "\n",
    "mod = get_test_model()\n",
    "custom_pi = CustomPI()\n",
    "# seq = transform.Sequential([relay.transform.AnnotateSpans(),\n",
    "#                             relay.transform.ToANormalForm(),\n",
    "#                             relay.transform.InferType()])\n",
    "\n",
    "with PassContext(instruments=[custom_pi]):\n",
    "    # mod = seq(mod)\n",
    "    mod = tvm.relay.transform.AnnotateSpans()(mod)\n",
    "    mod = tvm.relay.transform.ToANormalForm()(mod)\n",
    "    mod = tvm.relay.transform.InferType()(mod)\n",
    "\n",
    "assert \"InferType\" == \"\".join(custom_pi.events)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@pass_instrument\n",
    "class SkipPass:\n",
    "    def __init__(self, skip_pass_name):\n",
    "        self.skip_pass_name = skip_pass_name\n",
    "\n",
    "    def should_run(self, mod, info):\n",
    "        if self.skip_pass_name in info.name:\n",
    "            return False\n",
    "        return True\n",
    "\n",
    "skip_annotate = SkipPass(\"AnnotateSpans\")\n",
    "skip_anf = SkipPass(\"ToANormalForm\")\n",
    "\n",
    "@pass_instrument\n",
    "class PrintPassName:\n",
    "    def __init__(self):\n",
    "        self.events = []\n",
    "\n",
    "    def run_before_pass(self, mod, info):\n",
    "        self.events.append(info.name)\n",
    "\n",
    "mod = get_test_model()\n",
    "print_pass_name = PrintPassName()\n",
    "with tvm.transform.PassContext(instruments=[skip_annotate, skip_anf, print_pass_name]):\n",
    "    mod = tvm.relay.transform.AnnotateSpans()(mod)\n",
    "    mod = tvm.relay.transform.ToANormalForm()(mod)\n",
    "    mod = tvm.relay.transform.InferType()(mod)\n",
    "\n",
    "assert \"InferType\" == \"\".join(print_pass_name.events)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@pass_instrument\n",
    "class PassesCounter:\n",
    "    def __init__(self):\n",
    "        self.run_before_count = 0\n",
    "        self.run_after_count = 0\n",
    "\n",
    "    def __clear(self):\n",
    "        self.run_before_count = 0\n",
    "        self.run_after_count = 0\n",
    "\n",
    "    def enter_pass_ctx(self):\n",
    "        self.__clear()\n",
    "\n",
    "    def exit_pass_ctx(self):\n",
    "        self.__clear()\n",
    "\n",
    "    def run_before_pass(self, mod, info):\n",
    "        self.run_before_count = self.run_before_count + 1\n",
    "\n",
    "    def run_after_pass(self, mod, info):\n",
    "        self.run_after_count = self.run_after_count + 1\n",
    "\n",
    "mod = get_test_model()\n",
    "passes_counter = PassesCounter()\n",
    "with tvm.transform.PassContext(instruments=[passes_counter]):\n",
    "    tvm.relay.build(mod, \"llvm\")\n",
    "    assert passes_counter.run_after_count != 0\n",
    "    assert passes_counter.run_after_count == passes_counter.run_before_count\n",
    "\n",
    "# Out of pass context scope, should be reset\n",
    "assert passes_counter.run_before_count == 0\n",
    "assert passes_counter.run_after_count == 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "configs = PassContext.list_configs()\n",
    "\n",
    "assert len(configs) > 0\n",
    "assert \"relay.backend.use_auto_scheduler\" in configs.keys()\n",
    "assert configs[\"relay.backend.use_auto_scheduler\"][\"type\"] == \"IntImm\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "events = []\n",
    "\n",
    "@pass_instrument\n",
    "class PI:\n",
    "    def __init__(self, id):\n",
    "        self.id = id\n",
    "\n",
    "    def enter_pass_ctx(self):\n",
    "        events.append(self.id + \" enter_pass_ctx\")\n",
    "\n",
    "    def exit_pass_ctx(self):\n",
    "        events.append(self.id + \" exit_pass_ctx\")\n",
    "\n",
    "    def should_run(self, mod, info):\n",
    "        events.append(\"  \" + self.id + \" should_run\")\n",
    "        return True\n",
    "\n",
    "    def run_before_pass(self, mod, info):\n",
    "        events.append(\"  \" + self.id + \" run_before_pass\")\n",
    "\n",
    "    def run_after_pass(self, mod, info):\n",
    "        events.append(\"  \" + self.id + \" run_after_pass\")\n",
    "\n",
    "@tvm.transform.module_pass(opt_level=2)\n",
    "def transform1(mod, ctx):\n",
    "    events.append(\"    transform1 pass\")\n",
    "    return mod\n",
    "\n",
    "@tvm.transform.module_pass(opt_level=2)\n",
    "def transform2(mod, ctx):\n",
    "    events.append(\"    transform2 pass\")\n",
    "    return mod\n",
    "\n",
    "mod = get_test_model()\n",
    "with PassContext(instruments=[PI(\"%1\"), PI(\"%2\")]):\n",
    "    mod = transform1(mod)\n",
    "    mod = transform2(mod)\n",
    "\n",
    "assert (\n",
    "    \"%1 enter_pass_ctx\"\n",
    "    \"%2 enter_pass_ctx\"\n",
    "    \"  %1 should_run\"\n",
    "    \"  %2 should_run\"\n",
    "    \"  %1 run_before_pass\"\n",
    "    \"  %2 run_before_pass\"\n",
    "    \"    transform1 pass\"\n",
    "    \"  %1 run_after_pass\"\n",
    "    \"  %2 run_after_pass\"\n",
    "    \"  %1 should_run\"\n",
    "    \"  %2 should_run\"\n",
    "    \"  %1 run_before_pass\"\n",
    "    \"  %2 run_before_pass\"\n",
    "    \"    transform2 pass\"\n",
    "    \"  %1 run_after_pass\"\n",
    "    \"  %2 run_after_pass\"\n",
    "    \"%1 exit_pass_ctx\"\n",
    "    \"%2 exit_pass_ctx\" == \"\".join(events)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pass 去函数化"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tvm.relay.backend.interpreter import ConstructorValue\n",
    "from tvm.relay import transform, ExprVisitor, TypeVisitor\n",
    "from tvm.relay.testing import Prelude"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def has_func_type(t):\n",
    "    \"\"\"确定类型 t 是 FuncType 还是嵌套的 FuncType\"\"\"\n",
    "    class FuncTypeVisitor(TypeVisitor):\n",
    "        def __init__(self):\n",
    "            super().__init__()\n",
    "            self.has_func = False\n",
    "\n",
    "        def visit_func_type(self, ftt):\n",
    "            self.has_func = True\n",
    "\n",
    "    ftvisitor = FuncTypeVisitor()\n",
    "    ftvisitor.visit(t)\n",
    "    return ftvisitor.has_func"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "确定程序是否有高阶函数，高阶函数定义为:\n",
    "\n",
    "- 具有函数类型参数\n",
    "- 返回函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def assert_no_higher_order_functions(expr, mod):\n",
    "    class CheckFirstOrderVisitor(ExprVisitor):\n",
    "        def __init__(self, mod):\n",
    "            super().__init__()\n",
    "            self.mod = mod\n",
    "            self.hof = []\n",
    "            self.visited_gv = set()\n",
    "\n",
    "        def visit_call(self, call):\n",
    "            is_higher_order = False\n",
    "            # check return type\n",
    "            if has_func_type(call.checked_type):\n",
    "                is_higher_order = True\n",
    "            # check argument types\n",
    "            for a in call.args:\n",
    "                if has_func_type(a.checked_type):\n",
    "                    is_higher_order = True\n",
    "            # if it is higher order, save it for debugging later\n",
    "            if is_higher_order:\n",
    "                self.hof.append(call)\n",
    "            super().visit_call(call)\n",
    "\n",
    "        def visit_global_var(self, gv):\n",
    "            # visit global vars to visit entire program\n",
    "            if gv not in self.visited_gv:\n",
    "                self.visited_gv.add(gv)\n",
    "                self.visit(self.mod[gv])\n",
    "\n",
    "    mod = transform.InferType()(mod)\n",
    "    check_fo_visitor = CheckFirstOrderVisitor(mod)\n",
    "    check_fo_visitor.visit(expr)\n",
    "\n",
    "    nl = \"\\n--------\\n\"\n",
    "    errmsg = f\"\"\"found {len(check_fo_visitor.hof)} higher order functions:\n",
    "  {nl.join(expr.astext() for expr in check_fo_visitor.hof)}\"\"\"\n",
    "\n",
    "    assert len(check_fo_visitor.hof) == 0, errmsg"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "断言程序是去函数化的，并返回去函数化的模块，假设程序从 `mod['main']` 开始:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def defunctionalized(mod):\n",
    "    mod = transform.InferType()(mod)\n",
    "    mod[\"main\"] = transform.Defunctionalization(mod[\"main\"], mod)\n",
    "    mod = transform.InferType()(mod)\n",
    "    assert_no_higher_order_functions(mod[\"main\"], mod)\n",
    "    return mod"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# adt list to python list\n",
    "def to_list(mod, l):\n",
    "    list = mod.get_global_type_var(\"List\")\n",
    "    list_adt = mod[list]\n",
    "    cons = list_adt.constructors[0]\n",
    "    nil = list_adt.constructors[1]\n",
    "\n",
    "    assert isinstance(l, ConstructorValue)\n",
    "    val = l\n",
    "    ret = []\n",
    "    while True:\n",
    "        if val.tag == cons.tag:\n",
    "            ret.append(val.fields[0].numpy())\n",
    "            val = val.fields[1]\n",
    "        else:\n",
    "            assert val.tag == nil.tag\n",
    "            break\n",
    "    return ret\n",
    "\n",
    "\n",
    "# list to adt list\n",
    "def to_adt_list(mod, arr):\n",
    "    expr = mod[\"main\"]\n",
    "    l = mod.get_global_type_var(\"List\")\n",
    "    list_adt = mod[l]\n",
    "    cons = list_adt.constructors[0]\n",
    "    nil = list_adt.constructors[1]\n",
    "\n",
    "    li = nil()\n",
    "    for a in arr:\n",
    "        li = cons(relay.const(a), li)\n",
    "    adt = relay.create_executor(mod=mod).evaluate(li)\n",
    "    mod[\"main\"] = expr\n",
    "    return adt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tvm\n",
    "from tvm import relay\n",
    "import numpy as np\n",
    "\n",
    "code = \"\"\"\n",
    "#[version = \"0.0.5\"]\n",
    "def @simple[A, B](%f: fn(A) -> B, %xs: A) -> B {\n",
    "  %f(%xs)\n",
    "}\n",
    "def @main(%l: Tensor[(5, 5), float32]) -> Tensor[(5, 5), float32] {\n",
    "  %0 = fn[A](%x: A) -> A {\n",
    "    %x\n",
    "  };\n",
    "  @simple(%0, %l)\n",
    "}\n",
    "\"\"\"\n",
    "mod = tvm.parser.fromtext(code)\n",
    "defunc_mod = defunctionalized(mod)\n",
    "\n",
    "input = np.random.rand(5, 5).astype(\"float32\")\n",
    "out = relay.create_executor(\"debug\", mod=mod).evaluate()(input)\n",
    "defunc_out = relay.create_executor(\"debug\", mod=defunc_mod).evaluate()(input)\n",
    "np.testing.assert_equal(out.numpy(), defunc_out.numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "code = \"\"\"\n",
    "#[version = \"0.0.5\"]\n",
    "type List[A] {\n",
    "  Cons(A, List[A]),\n",
    "  Nil,\n",
    "}\n",
    "def @id[A](%x: A) -> A {\n",
    "  %x\n",
    "}\n",
    "def @map[A, B](%f: fn(A) -> B, %xs: List[A]) -> List[B] {\n",
    "  match (%xs) {\n",
    "    Cons(%x, %rest) => Cons(%f(%x), @map(%f, %rest)),\n",
    "    Nil => Nil,\n",
    "  }\n",
    "}\n",
    "def @main(%l: List[float32]) -> List[float32] {\n",
    "  @map(@id, %l)\n",
    "}\n",
    "\"\"\"\n",
    "mod = tvm.parser.fromtext(code)\n",
    "defunc_mod = defunctionalized(mod)\n",
    "\n",
    "input = np.random.rand(10).astype(\"float32\")\n",
    "\n",
    "out = relay.create_executor(\"debug\", mod=mod).evaluate(mod[\"main\"])(to_adt_list(mod, input))\n",
    "\n",
    "defunc_out = relay.create_executor(\"debug\", mod=defunc_mod).evaluate()(\n",
    "    to_adt_list(defunc_mod, input)\n",
    ")\n",
    "\n",
    "np.testing.assert_array_equal(to_list(mod, out), to_list(defunc_mod, defunc_out))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_recursive_datatype():\n",
    "    # CPS will create recursive datatype\n",
    "    code = \"\"\"\n",
    "#[version = \"0.0.5\"]\n",
    "type List[A] {\n",
    "  Cons(A, List[A]),\n",
    "  Nil,\n",
    "}\n",
    "def @sum(%f: fn(int32) -> int32, %k: List[int32]) -> int32 {\n",
    "  match (%k) {\n",
    "    Cons(%x, %rest) => %0 = fn(%n) {\n",
    "      %x + %f(%n)\n",
    "    };\n",
    "    @sum(%0, %rest),\n",
    "    Nil => %f(0),\n",
    "  }\n",
    "}\n",
    "def @id[A](%x: A) -> A {\n",
    "  %x\n",
    "}\n",
    "def @main(%l: List[int32]) -> int32 {\n",
    "  @sum(@id, %l)\n",
    "}\n",
    "\"\"\"\n",
    "    mod = tvm.parser.fromtext(code)\n",
    "    defunc_mod = defunctionalized(mod)\n",
    "\n",
    "    input = np.random.randint(1, 100, 10)\n",
    "\n",
    "    out = relay.create_executor(\"debug\", mod=mod).evaluate(mod[\"main\"])(to_adt_list(mod, input))\n",
    "\n",
    "    defunc_out = relay.create_executor(\"debug\", mod=defunc_mod).evaluate()(\n",
    "        to_adt_list(defunc_mod, input)\n",
    "    )\n",
    "\n",
    "    tvm.testing.assert_allclose(out.numpy(), defunc_out.numpy())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.13 ('py38': conda)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "28558e8daad512806f5c536a1a04c119185f99f65b79002708a12162d02a79c7"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
